import pandas as pd
from matplotlib import pyplot as plt
from chapter_2_3_sea_turtle import turtle_strategy



stock_turtle = turtle_strategy('002419', '20210101', '20221231')
print('-------------stock_turtle-------------')
print(stock_turtle)
initial_cash = 20000
positions = pd.DataFrame(index=stock_turtle.index).fillna(0.0)
# 每次交易为1手，即100股，仓位即买单和卖单的累加和
positions['stock'] = 100 * stock_turtle['order'].cumsum()
# 持仓价格
portfolio = positions.multiply(stock_turtle['price'], axis=0)
# 持仓市值 持仓数*价格
portfolio['holding_values'] = (positions.multiply(stock_turtle['price'], axis=0))

# 计算仓位变化
pos_diff = positions.diff()
# 剩余的现金 = 初始资金 - 仓位变化pos_diff*价格的累加和
portfolio['cash'] = initial_cash - (pos_diff.multiply(stock_turtle['price'], axis=0)).cumsum()
# 总资产即持仓股票市值加剩余现金
portfolio['total'] = portfolio['cash'] + portfolio['holding_values']
# stock_turtle.to_excel('dahua_002419_stock_turtle.xls')

# 可视化
plt.figure(figsize=(10, 5))
plt.plot(portfolio['total'])  # 总资产
plt.plot(portfolio['holding_values'], '--')  # 持仓市值
plt.grid()
plt.legend()
plt.show()
